import os, sys
import os.path as osp
import shutil
import numpy as np
import cv2

# 递归获取文件夹下所有符合条件的文件路径.
def get_all_filepath(root_dir_path, extension_tag=[], recursive = True):
    """
    Args:
        root_dir_path: 文件根目录.
        extension_tag: 文件后缀名组成的列表, Example: ["jpg", "png", "bmp"]
        如果为空 则获取所有
    Returns:
        特定文件夹下所有符合条件的文件路径.
    """
    def check_endswith(new_path, extension_tag):
        if len(extension_tag) == 0:
            return True
        for s in extension_tag:
            if new_path.lower().endswith(s):
                return True
        return False

    paths = []
    for f in os.listdir(root_dir_path):
        new_path = root_dir_path + os.sep + f
        if os.path.isfile(new_path) and check_endswith(new_path,  extension_tag):
            paths.append(new_path)
        elif os.path.isdir(new_path) and recursive:
            temp_list = get_all_filepath(new_path, extension_tag)
            paths.extend(temp_list)
    return paths

# 递归获取文件夹下所有最后一层子文件夹.
def get_last_dir(root_dir_path, recursive = True):
    """
    Args:
        root_dir_path: 文件夹根目录.
    Returns:
        特定文件夹下所有的文件夹路径.
    """
    paths = []
    for f in os.listdir(root_dir_path):
        new_path = root_dir_path + os.sep + f
        if osp.isdir(new_path):
            ps = get_last_dir(new_path, recursive)
            paths.extend(ps)
    if len(paths) == 0:
        paths.append(root_dir_path)
    return paths



def remove_image_by_mean(data_dir):
    last_dirs = get_last_dir(data_dir)

    for last_dir in last_dirs:
        files = get_all_filepath(last_dir, [".jpg"])
        mean_dict = {}

        for f in files:
            img = cv2.imread(f)
            mean = np.mean(img)
            mean = int(mean * 1e4)
            if mean_dict.get(mean, None): continue
            mean_dict[mean] = True
            mean_dict[f] = True

        for f in files:
            if mean_dict.get(f, None): continue
            os.remove(f)


if __name__ == "__main__":
    data_dir = sys.argv[1]
    remove_image_by_mean(data_dir)
